import math
import random
import struct
from typing import Any


def insert_random_character(s: str) -> str:
    """
    向 s 中下标为 pos 的位置插入一个随机 byte
    pos 为随机生成，范围为 [0, len(s)]
    插入的 byte 为随机生成，范围为 [32, 127]
    """
    # TODO
    pos = random.randint(0, len(s))
    randomByte = chr(random.randint(32, 127))
    return s[:pos] + randomByte + s[pos:]


def flip_random_bits(s: str) -> str:
    """
    基于 AFL 变异算法策略中的 bitflip 与 random havoc 实现相邻 N 位翻转（N = 1, 2, 4），其中 N 为随机生成
    从 s 中随机挑选一个 bit，将其与其后面 N - 1 位翻转（翻转即 0 -> 1; 1 -> 0）
    注意：不要越界
    """
    # TODO
    if not s:
        return s

    # 将字符串转换为字节数组（使用 utf-8 编码）
    try:
        byte_array = bytearray(s.encode('utf-8'))
    except UnicodeEncodeError:
        return s  # 无法编码的字符串跳过

    N = random.choice([1, 2, 4])
    total_bits = len(byte_array) * 8
    if total_bits < N:
        return s

    bit_pos = random.randint(0, total_bits - N)

    for i in range(N):
        current_bit = bit_pos + i
        byte_index = current_bit // 8
        bit_index = current_bit % 8
        byte_array[byte_index] ^= (1 << bit_index)

    # 将字节数组解码为原始编码，如果失败返回原始字符串
    try:
        return byte_array.decode('utf-8')
    except UnicodeDecodeError:
        return s  # 如果变异后无法解码，返回原字符串


def arithmetic_random_bytes(s: str) -> str:
    """
    基于 AFL 变异算法策略中的 arithmetic inc/dec 与 random havoc 实现相邻 N 字节随机增减（N = 1, 2, 4），其中 N 为随机生成
    字节随机增减：
        1. 取其中一个 byte，将其转换为数字 num1；
        2. 将 num1 加上一个 [-35, 35] 的随机数，得到 num2；
        3. 用 num2 所表示的 byte 替换该 byte
    从 s 中随机挑选一个 byte，将其与其后面 N - 1 个 bytes 进行字节随机增减
    注意：不要越界；如果出现单个字节在添加随机数之后，可以通过取模操作使该字节落在 [0, 255] 之间
    """
    # TODO
    if not s:
        return s

    # 编码为 UTF-8 字节数组
    try:
        byte_array = bytearray(s.encode('utf-8'))
    except UnicodeEncodeError:
        return s

    N = random.choice([1, 2, 4])
    if len(byte_array) < N:
        return s

    pos = random.randint(0, len(byte_array) - N)

    for i in range(N):
        original = byte_array[pos + i]
        delta = random.randint(-35, 35)
        byte_array[pos + i] = (original + delta) % 256

    # 解码为 UTF-8 字符串
    try:
        return byte_array.decode('utf-8')
    except UnicodeDecodeError:
        return s


def interesting_random_bytes(s: str) -> str:
    """
    基于 AFL 变异算法策略中的 interesting values 与 random havoc 实现相邻 N 字节随机替换为 interesting_value（N = 1, 2, 4），其中 N 为随机生成
    interesting_value 替换：
        1. 构建分别针对于 1, 2, 4 bytes 的 interesting_value 数组；
        2. 随机挑选 s 中相邻连续的 1, 2, 4 bytes，将其替换为相应 interesting_value 数组中的随机元素；
    注意：不要越界
    """
    # TODO
    if not s:
        return s

    interesting_values = {
        1: [0x00, 0xFF, 0x7F, 0x80],
        2: [0x0000, 0xFFFF, 0x7FFF, 0x8000],
        4: [0x00000000, 0xFFFFFFFF, 0x7FFFFFFF, 0x80000000],
    }

    try:
        byte_array = bytearray(s.encode('utf-8'))
    except UnicodeEncodeError:
        return s

    N = random.choice([1, 2, 4])
    if len(byte_array) < N:
        return s

    pos = random.randint(0, len(byte_array) - N)
    val = random.choice(interesting_values[N])
    packed = val.to_bytes(N, byteorder='little', signed=False)
    byte_array[pos:pos + N] = packed

    try:
        return byte_array.decode('utf-8')
    except UnicodeDecodeError:
        return s


def havoc_random_insert(s: str):
    """
    基于 AFL 变异算法策略中的 random havoc 实现随机插入
    随机选取一个位置，插入一段的内容，其中 75% 的概率是插入原文中的任意一段随机长度的内容，25% 的概率是插入一段随机长度的 bytes
    """
    # TODO
    pos = random.randint(0, len(s))
    choice = random.random()
    if choice < 0.75 and len(s) > 0:
        # 插入原文中的一段
        start = random.randint(0, len(s) - 1)
        end = random.randint(start + 1, len(s))
        insertData = s[start:end]
    else:
        # 插入随机字符
        length = random.randint(1, 8)
        insertData = ''.join([chr(random.randint(32, 127)) for _ in range(length)])
    return s[:pos] + insertData + s[pos:]


def havoc_random_replace(s: str):
    """
    基于 AFL 变异算法策略中的 random havoc 实现随机替换
    随机选取一个位置，替换随后一段随机长度的内容，其中 75% 的概率是替换为原文中的任意一段随机长度的内容，25% 的概率是替换为一段随机长度的 bytes
    """
    # TODO
    if not s:
        return s
    pos = random.randint(0, len(s) - 1)
    maxReplaceLen = len(s) - pos
    length = random.randint(1, maxReplaceLen)
    choice = random.random()
    if choice < 0.75 and len(s) > length:
        # 替换为原文中的一段
        start = random.randint(0, len(s) - length)
        replaceData = s[start:start + length]
    else:
        # 替换为随机字符
        replaceData = ''.join([chr(random.randint(32, 127)) for _ in range(length)])
    return s[:pos] + replaceData + s[pos + length:]


def duplicate_random_chunk(s: str) -> str:
    """
    随机复制原始字符串中的一段内容，并插入到随机位置。
    """
    if not s:
        return s
    start = random.randint(0, len(s) - 1)
    end = random.randint(start + 1, len(s))
    chunk = s[start:end]
    insertPos = random.randint(0, len(s))
    return s[:insertPos] + chunk + s[insertPos:]


def test_all_mutators(inputStr: str):
    print(f"原始输入: {repr(inputStr)}\n" + "-" * 50)
    mutators = [
        ("insert_random_character", insert_random_character),
        ("flip_random_bits", flip_random_bits),
        ("arithmetic_random_bytes", arithmetic_random_bytes),
        ("interesting_random_bytes", interesting_random_bytes),
        ("havoc_random_insert", havoc_random_insert),
        ("havoc_random_replace", havoc_random_replace),
        ("duplicate_random_chunk", duplicate_random_chunk)
    ]

    for name, func in mutators:
        try:
            mutated = func(inputStr)
            print(f"{name} -> {repr(mutated)}")
        except Exception as e:
            print(f"{name} 抛出异常: {e}")


class Mutator:

    def __init__(self) -> None:
        """Constructor"""
        self.mutators = [
            insert_random_character,
            flip_random_bits,
            arithmetic_random_bytes,
            interesting_random_bytes,
            havoc_random_insert,
            havoc_random_replace,
            duplicate_random_chunk  # 实现的额外 Mutator
        ]

    def mutate(self, inp: Any) -> Any:
        mutator = random.choice(self.mutators)
        return mutator(inp)


if __name__ == "__main__":
    testInput = "id=1234&user=test&flag=0"
    test_all_mutators(testInput)
